#include <stdio.h>
#include <stdlib.h>

typedef struct tree_node
{
    int data;
    // 二叉平衡树 AVL中需要使用该成员,确保树在增删过程中不被劣化
    int height;
    struct tree_node *left, *right;
} TreeNode;

// AVL 树结构体
typedef struct
{
    TreeNode *root;
} AVLTree;

// “节点高度”是指从该节点到它的最远叶节点的距离，即所经过的“边”的数量。
// 注意: 叶节点的高度为 0, 而空节点的高度为 -1

//初始化AVL树节点
TreeNode *new_tree_node(int data)
{
    TreeNode *n = malloc(sizeof(TreeNode));
    n->data = data;
    n->height = 0;
    n->left = NULL;
    n->right = NULL;

    return n;
}

// 我们将创建两个工具函数，分别用于获取和更新节点的高度：

//获取新节点的高度
int height(TreeNode *node)
{
    if (node != NULL)
    {
        return node->height;
    }
    return -1;
}

// 更新节点高度
void updateHeight(TreeNode *node)
{
    int lh = height(node->left);   //该节点的左子树节点的高度
    int rh = height(node->right);  //该节点的右子树节点的高度
    node->height = lh > rh ? lh + 1 : rh + 1;
}

// 平衡因子,判断当前节点左右子树高度差
// 同时规定空节点的平衡因子为 0
// 设平衡因子为 f ，则一棵 AVL 树的任意节点的平衡因子皆满足 -1< f <1。
int balance_factor(TreeNode *node)
{
    if (node == NULL)
    {
        return 0;
    }
    return height(node->left) - height(node->right); //(平衡因子 = 左子树高度 - 右子树高度)
}

// 相关的旋转方法
// 右旋
TreeNode *right_rotate(TreeNode *n)
{
    TreeNode *child = n->left;
    TreeNode *grand_child = child->right;
    child->right = n;
    n->left = grand_child;
    updateHeight(n);     // 先计算n的高度
    updateHeight(child); // 再计算child高度
    // 因为右旋后child变成父节点了
    return child;
}

// 左旋转
TreeNode *left_rotate(TreeNode *n)
{
    TreeNode *child = n->right;
    TreeNode *grand_child = child->left;
    child->left = n;
    n->right = grand_child;
    updateHeight(n);     // 先计算n的高度
    updateHeight(child); // 再计算child高度
    return child;
}

// 总的旋转操作,会根据平衡因子分析当前失衡节点该如何选择
// 对于左子树失衡(左偏树),并且子节点平衡因子 >= 0 进行右旋
// 对于左子树失衡(左偏树),并且子节点平衡因子 < 0 进行先左旋后右旋
// 对于右子树失衡(右偏树),并且子节点平衡因子 <= 0 进行左旋
// 对于右子树失衡(右偏树),并且子节点平衡因子 > 0 进行先右旋后左旋
TreeNode *rotate(TreeNode *n)
{
    int bf = balance_factor(n);
    // 左偏树
    if (bf > 1)
    { // 右旋
        if (balance_factor(n->left) >= 0)
        {
            return right_rotate(n);
        }
        else
        {
            // 先左旋再右旋
            n->left = left_rotate(n->left);
            return right_rotate(n);
        }
    }
    // 右偏树
    else if (bf < -1)
    {
        // 左旋
        if (balance_factor(n->right) <= 0) 
        {
            return left_rotate(n);
        }
        else
        {
            n->right = right_rotate(n->right);
            return left_rotate(n);
        }
    }

    return n;
}
// 插入节点辅助函数,目的方便递归
TreeNode *insert_helper(TreeNode *node, int data)
{
    // 若AVL二叉树为空,第一次插入的节点直接作为根节点
    if (node == NULL)
    {
        return new_tree_node(data);
    }

    // 插入, AVL满足二叉搜索树的条件

    if (data > node->data)
    {
        node->right = insert_helper(node->right, data);
    }
    else if (data < node->data)
    {
        node->left = insert_helper(node->left, data);
    }
    else
    {
        return node;
    }
    // 更新节点的高度
    updateHeight(node);

    // 测试(该节点是否是失衡节点)并旋转当前节点
    node = rotate(node);

    return node;
}

// 插入节点
void insert(AVLTree *tree, int data)
{
    // 每次插入新节点都要重新计算节点高度,以及是否平衡
    // 若不平衡就需要旋转,旋转的过程中AVL树的根节点可能会改变
    tree->root = insert_helper(tree->root, data);
}

TreeNode *remove_helper(TreeNode *node, int data)
{
    if (node == NULL)
    {
        return NULL;
    }

    TreeNode *child, *grand_child;

    if (node->data > data)
    {
        node->left = remove_helper(node->left, data);
    }
    else if (node->data < data)
    {
        node->right = remove_helper(node->right, data);
    }
    else // 找到目标元素
    {
        if (node->left == NULL || node->right == NULL)
        {
            child = node->left ? node->left : node->right;
            free(node);
            if (child)
            {
                node = child;
            }
            else
            {
                return NULL;
            }
        }
        else
        {
            // 删除二叉搜索树度为2的节点,需要中序遍历找到
            TreeNode *tmp = node->right;
            while (tmp->left != NULL)
            {
                tmp = tmp->left;
            }
            int tmp_val = tmp->data;
            node->right = remove_helper(node->right, tmp_val);
            node->data = tmp_val;
        }
    }
    updateHeight(node);
    node = rotate(node);
    return node;
}

// 删除节点
void remove_item(AVLTree *tree, int data)
{
    tree->root = remove_helper(tree->root, data);
}

int main(int argc, char const *argv[])
{
    AVLTree *tree = malloc(sizeof(AVLTree));
    tree->root = NULL;

    insert(tree, 18);
    printf("%d \n", tree->root->data);

    insert(tree, 15);
    printf("%d \n", tree->root->data);

    insert(tree, 9);
    insert(tree, 21);
    insert(tree, 17);
    // printf("%d \n", tree->root->data);

    // printf("%d \n", tree->root->right->data);

    // printf("%d \n", tree->root->right->height);
    remove_item(tree, 15);
    printf("%d \n", tree->root->data);

    remove_item(tree, 9);
    printf("%d \n", tree->root->data);

    printf("%d \n", tree->root->left->data);

    return 0;
}
